import mxnet as mx
import numpy as np
import mxnet.ndarray as nd
import mxnet.gluon as gluon
import mxnet.gluon.nn as nn
from models.defomableconvnets.config import update_config,config
from models.defomableconvnets.resnet_v1_101_fpn_dcn_rcnn import resnet_v1_101_fpn_dcn_rcnn
from models.defomableconvnets.image import transform
from models.defomableconvnets.bbox.bbox_transform import bbox_pred,bbox_transform,clip_boxes
from models.defomableconvnets.nms.nms import gpu_nms_wrapper
from utils.block import SymbolBlock
from utils.common import lsdir
import cv2
def get_symbol(cfg):
    sym = resnet_v1_101_fpn_dcn_rcnn().get_symbol(cfg,is_train=False)
    return sym

def im_detect_bbox_aug(net, nms_wrapper, img_path, scales, pixel_means, ctx, threshold = 1e-3, viz = False):

    all_bboxes = []
    all_scores = []
    img_ori = cv2.imread(img_path)
    for scale_min, scale_max in scales:
        fscale = 1.0 * scale_min / min(img_ori.shape[:2])
        img_resized = cv2.resize(img_ori,(0,0),fx = fscale,fy = fscale)
        h,w,c = img_resized.shape
        h_padded = h if h %32==0 else h + 32 - h % 32
        w_padded = w if w %32==0 else w + 32 - w % 32
        img_padded = np.zeros(shape = (h_padded,w_padded,c), dtype = img_resized.dtype)
        img_padded[:h,:w,:] = img_resized
        img = transform(img_padded,pixel_means=pixel_means)
        im_info = nd.array([[h_padded,w_padded,1.0]], ctx= ctx[0])
        data = nd.array(img,ctx=ctx[0])

        rois, scores, bbox_deltas = net(data,im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])
        bbox /= fscale
        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

        # flip
        rois, scores, bbox_deltas = net(data[:,:,:,::-1],im_info)
        rois = rois[:,1:].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas[0].asnumpy())
        bbox = clip_boxes(bbox,data.shape[2:4])

        tmp = bbox[:, 0::4].copy()
        bbox[:, 0::4] =data.shape[3] - bbox[:, 2::4] #x0 = w - x0
        bbox[:, 2::4] =data.shape[3] - tmp # x1 = w -x1
        bbox /= fscale

        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

    all_bboxes = np.concatenate(all_bboxes,axis = 0)
    all_scores = np.concatenate(all_scores,axis = 0)
    pred_bboxes = []
    pred_scores = []
    pred_clsid = []
    for j in range(1, all_scores.shape[1]):
        cls_scores = all_scores[:, j, np.newaxis]
        cls_boxes = all_bboxes[:, 4:8] if config.CLASS_AGNOSTIC else all_bboxes[:, j * 4:(j + 1) * 4]
        cls_dets = np.hstack((cls_boxes, cls_scores))
        keep = nms_wrapper(cls_dets.astype('f'))
        cls_dets = cls_dets[keep, :]
        cls_dets = cls_dets[cls_dets[:, -1] > threshold, :]
        pred_bboxes.append(cls_dets[:,:4])
        pred_scores.append(cls_dets[:,4])
        pred_clsid.append(j * np.ones(shape=(cls_dets.shape[0],), dtype=np.int))
    pred_bboxes = np.concatenate(pred_bboxes,axis = 0)
    pred_scores = np.concatenate(pred_scores,axis = 0)
    pred_clsid = np.concatenate(pred_clsid,axis = 0)
    if viz:
        import gluoncv
        import matplotlib.pyplot as plt
        gluoncv.utils.viz.plot_bbox(img_ori[:,:,::-1], bboxes=pred_bboxes,scores=pred_scores,labels=pred_clsid,thresh=.5)
        plt.show()
    return pred_bboxes,pred_scores,pred_clsid
if __name__ == '__main__':
    import os,tqdm,time
    os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
    gpu_id = 8
    ctx = [mx.gpu(gpu_id)]
    update_config("/home/zyx/Downloads/resnet_v1_101_coco_trainval_fpn_dcn_end2end_ohem.yaml")
    sym = get_symbol(config)
    net = SymbolBlock(sym= sym,input_names=["data","im_info"],
                      resume="/data1/zyx/yks/AiChallenger/bdd-100k/Deformable-ConvNets/output/fpn/"
                             "coco/resnet_v1_101_coco_trainval_fpn_dcn_end2end_ohem/retail/fpn_coco-0005.params")
    net.collect_params().reset_ctx(ctx)
    im_names = list(lsdir("/data1/zyx/yks/dataset/retail/test_images/test",suffix=".jpg"))
    nms_wrapper = nms = gpu_nms_wrapper(config.TEST.NMS, gpu_id)
    with open("output/retail_deformable_multi_scales_flipped_{}.txt".format(time.time()), "wt") as f:
        for im_name in tqdm.tqdm(im_names):
            bboxes, scores, labels = im_detect_bbox_aug(net,nms_wrapper,im_name, config.TEST_SCALES, config.network.PIXEL_MEANS,ctx=ctx)
            for bbox, score, label in zip(bboxes,scores,labels):
                f.write("{}\t{}\t{}\t{}\t{}\t{}\t{}\r\n".format(os.path.basename(im_name),
                                                                bbox[0],bbox[1],bbox[2],bbox[3],
                                                                score, label))